import os
import os.path as osp
import math
import pandas as pd
from PIL import Image
'''
 txt: 
    cls0 0.0
    x float_0-1
    y float_0-1
    w float_0-1
    h float_0-1
    angle int_degree

 xml:
    x float_3000
    y float_3000
    w float_3000
    h float_3000
    angle float_rad 
'''

BASEDIR = osp.dirname(osp.abspath(__file__))

cls_dict = {
    'A': '0.0',
    'B': '1.0',
    'C': '2.0',
    'D': '3.0',
    'E': '4.0',
    'F': '5.0',
    'G': '6.0',
    'H': '7.0',
    'I': '8.0',
    'J': '9.0',
    'K': '10.0',
}

inv_cls_dict = {}
for k in cls_dict.keys():
    inv_cls_dict[cls_dict[k]] = k
print(inv_cls_dict)


IMGDIR = osp.abspath(osp.join(BASEDIR, 'imgs'))
TXTDIR = osp.abspath(osp.join(BASEDIR, 'txts'))
OUTDIR = osp.abspath(osp.join(BASEDIR, 'xmls'))
if not osp.exists(OUTDIR):
    os.makedirs(OUTDIR)

all_img_names = [name[:-4] for name in os.listdir(IMGDIR) if name.endswith('.png')]
all_txt_names = [name[:-4] for name in os.listdir(TXTDIR) if name.endswith('.txt')]
intersection_names = set(all_img_names) & set(all_txt_names)

img_names = [name for name in os.listdir(IMGDIR) if name.endswith('.png')]
for img_name in img_names:
    img_path = osp.abspath(osp.join(IMGDIR, img_name))
    filename = img_path.split(os.sep)[-1].split('.')[0]

    if filename not in intersection_names:
    	continue

    folder = img_path.split(os.sep)[-2]

    txt_name = filename + '.txt'
    new_xml_name = filename + '.xml'

    Indentation = '  '

    IMGSIZE_W, IMGSIZE_H = Image.open(img_path).size

    with open(os.path.join(TXTDIR, txt_name)) as fp:
        lines = fp.readlines()
    # print(lines)

    box_info_list = [line.strip().split() for line in lines]
    # print(box_info_list)

    heads = ['cls0', 'x', 'y', 'w', 'h', 'angle']
    boxes_df = pd.DataFrame(box_info_list, columns=heads)
    # print(boxes_df)
    # print(len(boxes_df))


    # construct objects
    objs = []
    for i in range(len(boxes_df)):
        cls0 = boxes_df.loc[i]['cls0']
        x = eval(boxes_df.loc[i]['x'])
        y = eval(boxes_df.loc[i]['y'])
        w = eval(boxes_df.loc[i]['w'])
        h = eval(boxes_df.loc[i]['h'])
        angle = eval(boxes_df.loc[i]['angle'])

        # convert to xml value range
        x = str(x * IMGSIZE_W)
        y = str(y * IMGSIZE_H)
        w = str(w * IMGSIZE_W)
        h = str(h * IMGSIZE_H)
        angle = str(angle * math.pi / 180)

        # generate characters in xml
        one_obj = \
        Indentation + '<object>\n' + \
        Indentation*2 + '<type>robndbox</type>\n' + \
        Indentation*2 + f'<name>{inv_cls_dict[cls0]}</name>\n' + \
        Indentation*2 + '<pose>Unspecified</pose>\n' + \
        Indentation*2 + '<truncated>0</truncated>\n' + \
        Indentation*2 + '<difficult>0</difficult>\n' + \
        Indentation*2 + '<robndbox>\n' + \
        Indentation*3 + '<cx>' + x + '</cx>\n' + \
        Indentation*3 + '<cy>' + y + '</cy>\n' + \
        Indentation*3 + '<w>' + w + '</w>\n' + \
        Indentation*3 + '<h>' + h + '</h>\n' + \
        Indentation*3 + '<angle>' + angle + '</angle>\n' + \
        Indentation*2 + '</robndbox>\n' + \
        Indentation + '</object>\n'

        # print(one_obj, end='')
        # print(type(one_obj))
        objs.append(one_obj)


    # construct header and tail
    xml_header = \
    f'''<annotation verified="no">
      <folder>{folder}</folder>
      <filename>{filename}</filename>
      <path>{osp.basename(img_path)}</path>
      <source>
        <database>Unknown</database>
      </source>
      <size>
        <width>{IMGSIZE_W}</width>
        <height>{IMGSIZE_H}</height>
        <depth>3</depth>
      </size>
      <segmented>0</segmented>
    '''


    xml_tail = \
    '''</annotation>
    '''

    # write to xml file
    with open(os.path.join(OUTDIR, new_xml_name), 'w') as fp:
        fp.write(xml_header)
        for i in range(len(objs)):
            fp.write(objs[i])
        fp.write(xml_tail)
